#!/usr/bin/env python
# 车标分类算法

import os
MODELS_DIR = 'models/'
if not os.path.exists(MODELS_DIR):
    os.mkdir(MODELS_DIR)
MODEL_TF = MODELS_DIR + 'model'
MODEL_NO_QUANT_TFLITE = MODELS_DIR + 'model_no_quant.tflite'
MODEL_TFLITE = MODELS_DIR + 'model.tflite'
MODEL_TFLITE_MICRO = MODELS_DIR + 'model.cc'

import tensorflow as tf
from tensorflow import keras
import numpy as np
import matplotlib.pyplot as plt

fashion_mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

#print(train_images[0])
print("red car")
print(type(train_images), len(train_images))
#exit()
#print(train_labels)

#print(train_images.shape)

#plt.figure()
#plt.imshow(train_images[0])
#plt.colorbar()
#plt.grid(False)
#plt.show()

class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

train_images = train_images
test_images = test_images
train_images = np.array(train_images, dtype=np.float32)

#plt.figure(figsize=(10,10))
#for i in range(25):
#    plt.subplot(5,5,i+1)
#    plt.xticks([])
#    plt.yticks([])
#    plt.grid(False)
#    plt.imshow(train_images[i], cmap=plt.cm.binary)
#    plt.xlabel(class_names[train_labels[i]])
#plt.show()

model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10)
])

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
# 如果模型已经存在，加载整个模型
if os.path.exists(MODELS_DIR) and os.path.exists(MODEL_TF) and 0 != len(os.listdir(MODEL_TF)):
    try:
        #model.load_weights(MODEL_TF + 'saved_model.pb')
        model = keras.models.load_model(MODEL_TF)
    except:
        print("load weights failed")
        exit()
# 训练模型
else:
    print(f"{MODEL_TF} no exist ???")
    model.fit(train_images, train_labels, epochs=10)

# 查看  model 概要
print(model.summary())
# 测试模型
test_loss, test_acc = model.evaluate(test_images,  test_labels, verbose=2)

print('\nTest accuracy:', test_acc)

probability_model = tf.keras.Sequential([model,
                                         tf.keras.layers.Softmax()])
predictions = probability_model.predict(test_images)
print(type(test_images), test_images.shape, " red")

print(predictions[0])

# 测试自定义图片
from PIL import Image
import sys

if len(sys.argv) == 1:
    IMG_SRC = '../web_pics/boots.jpg'
else:
    IMG_SRC = '../web_pics/' + sys.argv[1]

IMG_TARGET = IMG_SRC.split('/')[-1].split('.')[0] + '_gry_test' + '.bmp'
CC_TARGET = IMG_SRC.split('/')[-1].split('.')[0] + '.cc'
print(CC_TARGET)
print(IMG_TARGET)

raw_image = Image.open(IMG_SRC)
image = raw_image.resize((28, 28))
image_gray = image.convert('L')
image_gray_array = np.array(image_gray, dtype=np.float32)
image_gray_array = image_gray_array.reshape(1,28,28)
#image_gray_array = image_gray_array / 255.0
image_gray_array = 255 - image_gray_array
image_gray_array.save(IMG_TARGET)
predictions = probability_model.predict(image_gray_array)
np.savetxt(CC_TARGET, image_gray_array.reshape(1, 28*28), fmt='%d', delimiter=',', header="const unsigned char g_person[]={", footer="};", comments='')

print("red prefictions")
print(predictions[0])

plt.figure()
plt.subplot(3,1,1)
plt.imshow(raw_image, cmap=plt.cm.binary)
plt.subplot(3,1,2)
plt.imshow(image_gray_array[0], cmap=plt.cm.binary)
plt.subplot(3,1,3)
plt.imshow(test_images[0], cmap=plt.cm.binary)
plt.show()
#exit()

#plt.figure(figsize=(10,10))
#plt.xticks([])
#plt.yticks([])
#plt.grid(False)
#plt.imshow(test_images[0], cmap=plt.cm.binary)
#plt.show()

# Save the model to disk
model.save(MODEL_TF)

# Convert the model to the TensorFlow Lite format without quantization
converter = tf.lite.TFLiteConverter.from_saved_model(MODEL_TF)
model_no_quant_tflite = converter.convert()

# Save the model to disk
open(MODEL_NO_QUANT_TFLITE, "wb").write(model_no_quant_tflite)

# Convert the model to the TensorFlow Lite format with quantization
def representative_dataset():
  for i in range(500):
    # 强制转换为 float32 类型
    array = np.array(train_images[i], dtype=np.float32)
    #print(array.shape, "red dbg")
    #print(type(train_images[0].dtype), type(array[0].dtype), type(array), len(array))
    yield([array])
  
# Set the optimization flag.
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# Enforce integer only quantization
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.float32
converter.inference_output_type = tf.float32
# Provide a representative dataset to ensure we quantize correctly.
converter.representative_dataset = representative_dataset
model_tflite = converter.convert()

# Save the model to disk
open(MODEL_TFLITE, "wb").write(model_tflite)
